import os
import pandas as pd
import backtrader as bt

pd.set_option('display.max_columns', None)
pd.set_option('display.max_rows', None)

dir_path = r'E:\mygit\learn_backtrader\Data'
daily_file = 'daily_price.csv'
trade_file = 'trade_info.csv'

daily_price = pd.read_csv(os.path.join(dir_path, daily_file), parse_dates=['datetime'])
trade_info = pd.read_csv(os.path.join(dir_path, trade_file), parse_dates=['trade_date'])

data = pd.DataFrame(daily_price['datetime'].unique(), columns=['datetime'])
# print(data)
df = daily_price.query('sec_code == "600497.SH"')[['datetime', 'open', 'high', 'low', 'close', 'volume', 'openinterest']]
# print(data)
# print(20*'-')
# print(df)
# print(20*'*')
# df2 = daily_price.query('sec_code == "600466.SH"')[['open', 'high', 'low', 'close', 'volume', 'openinterest']]
data_ = pd.merge(data, df, how='left', on='datetime')
df = df.set_index('datetime')
print(df)
print(20*'-')
print(data_)
data_ = data_.set_index("datetime")
# data_ = pd.merge(data, df2, left_index=True, right_index=True, how='left')
data_.loc[:, ['volume', 'openinterest']] = data_.loc[:, ['volume', 'openinterest']].fillna(0)
data_.loc[:, ['open', 'high', 'low', 'close']] = data_.loc[:, ['open', 'high', 'low', 'close']].fillna(method='pad')
data_.loc[:, ['open', 'high', 'low', 'close']] = data_.loc[:, ['open', 'high', 'low', 'close']].fillna(0)

buy_stock = trade_info
res = pd.to_datetime(buy_stock['trade_date'].unique()).tolist()
# print(res)
# print(daily_price.query('sec_code == "600497.SH"'))

# cerebro = bt.Cerebro()
# # 打印初始资金
# print('Starting Portfolio Value: %.2f' % cerebro.broker.getvalue())
# # 启动回测
# cerebro.run()
# # 打印回测完成后的资金
# print('Final Portfolio Value: %.2f' % cerebro.broker.getvalue())
